import pandas as pd
from sklearn.model_selection import train_test_split
import jieba


def write_data(data_df):

    # call_id = row.call_id
    content = data_df.context
    content = '。'.join([context[3:] for context in content.split('\n')])
    text = content.replace('\n', '。').replace(' ', '').replace('　', '')
    return text


def change_label(data_df, cls):
    label = data_df.label
    # if label in ['1-红色', '2-橙色', '3-黄色', '4-蓝色']:
    #     label = '1-红色'
    return cls[label]


if __name__ == '__main__':
    df = pd.read_excel('file/20230912.xlsx', engine='openpyxl', dtype='object')
    column = df.columns
    train_df_yes = pd.DataFrame(columns=column)
    test_df_yes = pd.DataFrame(columns=column)
    train_df_no = pd.DataFrame(columns=column)
    test_df_no = pd.DataFrame(columns=column)
    train_df = pd.DataFrame(columns=column)
    test_df = pd.DataFrame(columns=column)
    df = df.dropna()
    # df['label'] = df.apply(lambda x: change_label(x), axis=1)
    cls_yes = {'1-红色': 'yes', '2-橙色': 'yes', '3-黄色': 'yes'}
    cls_no = {'5-绿色': 'no', '4-蓝色': 'no'}
    cls_3 = {'6-内部通话': 'inner', '1-红色': 'yes', '2-橙色': 'yes', '3-黄色': 'yes', '4-蓝色': 'no', '5-绿色': 'no'}
    flag_cls2 = 0
    # cls = {'6-内部通话': 0, '1-红色': 1, '5-绿色': 2}
    if flag_cls2 == 1:
        for class_label, value in cls_yes.items():
            class_data = df[(df['label'] == class_label)]
            class_data['label'] = class_data.apply(lambda x: change_label(x, cls_yes), axis=1)
            train_part, test_part = train_test_split(class_data, test_size=0.2)
            train_df_yes = pd.concat([train_df_yes, train_part])
            test_df_yes = pd.concat([test_df_yes, test_part])
        for class_label, value in cls_no.items():
            class_data = df[(df['label'] == class_label)]
            class_data['label'] = class_data.apply(lambda x: change_label(x, cls_no), axis=1)
            train_part, test_part = train_test_split(class_data, test_size=0.2)
            train_df_no = pd.concat([train_df_no, train_part])
            test_df_no = pd.concat([test_df_no, test_part])
        train_df_no = train_df_no.sample(n=train_df_yes.shape[0])
        test_df_no = test_df_no.sample(n=test_df_yes.shape[0])
        train_df = pd.concat([train_df_yes, train_df_no])
        test_df = pd.concat([test_df_yes, test_df_no])
    else:
        for class_label, value in cls_3.items():
            class_data = df[(df['label'] == class_label)]
            class_data['label'] = class_data.apply(lambda x: change_label(x, cls_3), axis=1)
            train_part, test_part = train_test_split(class_data, test_size=0.2)
            train_df = pd.concat([train_df, train_part])
            test_df = pd.concat([test_df, test_part])
    train_df['context'] = df.apply(lambda x: write_data(x), axis=1)
    test_df['context'] = df.apply(lambda x: write_data(x), axis=1)
    train_df = train_df.sample(n=train_df.shape[0])
    test_df = test_df.sample(n=test_df.shape[0])
    train_df.to_csv('dataset/train_dataset_0912_3_noblue.csv', index=False)
    test_df.to_csv('dataset/valid_dataset_0912_3_noblue.csv', index=False)